"""
该文件仅用于将视频数据批量处理为训练数据集，核心代码与main.py一致。数据格式：角度1 角度2 角度3; 距离1 距离2 距离3; 3; 标签
Batch processing video data into training datasets. An example: ang1 ang2 ang3; dis1 dis2 dis3; 3; label
"""
import os
import sys
import cv2
import math
import config
import hand_extraction as he


# 路径全局变量
FILE_ROUTE = 'D:/Products/train/'
OUT_FILE = FILE_ROUTE + 'data.txt'
# 模型超参数
REVISE_WIDTH = config.REVISE_WIDTH
REVISE_HEIGHT = config.REVISE_HEIGHT
MIN_ANGLE = config.MIN_ANGLE
MIN_DISTANCE = config.MIN_DISTANCE
MAX_DISTANCE = config.MAX_DISTANCE


if __name__ == '__main__':

    outer = 0
    while True:
        inner = 1
        file_name = str(outer)
        is_file_exists = os.path.exists(FILE_ROUTE + file_name)
        if not is_file_exists:
            print('All videos have been manipulated!')
            sys.exit(0)
        while True:
            video_name = FILE_ROUTE + file_name + "/" + file_name + " (" + str(inner) + ").mp4"
            is_video_exists = os.path.exists(video_name)
            if not is_video_exists:
                print('No video called: ', video_name)
                break
            cen_point = []   # 轨迹质心
            key_point = []    # 轨迹关键点
            angle = []  # 临时角度
            angle_out = []  # 输出的角度
            re_angle_out = []   # 输出的逆序角度
            distance_out = []   # 输出的距离
            print('<------ Category %s  Sample %s ------>' % (file_name, inner))
            cap = cv2.VideoCapture(video_name)  # 获取视频流
            while True:
                ret, frame = cap.read()
                if frame is None:   # 判断视频流中的帧是否存在
                    break
                if frame.shape[1] > frame.shape[0]:
                    frame = cv2.transpose(frame)
                    frame = cv2.flip(frame, 1)
                frame = cv2.resize(frame, (REVISE_WIDTH, REVISE_HEIGHT))

                # 创建全黑图
                show_img = frame.copy()
                show_img[:, :] = 0

                # 获取手势质心
                center = he.get_centroid(frame)
                # cv2.circle(show_img, center, 4, (255, 0, 0), cv2.FILLED)

                # 获取关键点之间的角度与距离
                if len(cen_point) == 0:
                    key_point.append(center)
                elif len(cen_point) == 1:
                    angle.append(he.get_angle(cen_point[-1], center))
                else:
                    condition_ang = math.fabs(he.get_angle(cen_point[-1], center) - angle[-1]) > MIN_ANGLE
                    condition_dis = he.get_distance(center, key_point[-1]) > MAX_DISTANCE
                    if condition_ang or condition_dis:
                        if he.get_distance(center, key_point[-1]) > MIN_DISTANCE:
                            angle.append(he.get_angle(cen_point[-1], center))
                            key_point.append(center)
                            cur_angle = round(he.get_angle(key_point[-2], key_point[-1]) / (2 * math.pi), 4)
                            cur_distance = round(he.get_distance(key_point[-2], key_point[-1]), 2)
                            angle_out.append(cur_angle)
                            distance_out.append(cur_distance)
                            if outer == 20:
                                reverse_angle = round(he.get_angle(key_point[-1], key_point[-2]) / (2 * math.pi), 4)
                                re_angle_out.append(reverse_angle)
                            print('(%s,  %s)' % (cur_angle, cur_distance))
                cen_point.append(center)

                # 绘制轨迹和关键点
                for i in range(len(key_point)):
                    if i > 0:
                        cv2.line(show_img, key_point[i], key_point[i - 1], (255, 255, 255), 4)

                # 把轨迹存为图片
                cv2.imwrite('D:/Products/Pics/' + file_name + " (" + str(inner) + ").jpg", show_img)

            # 文件存储
            out_file = open(OUT_FILE, 'a+')
            index = 0
            for i in range(len(angle_out)):
                out_file.write(" " + str(angle_out[i]))
                index += 1
            out_file.write(";")
            for i in range(len(distance_out)):
                out_file.write(" " + str(distance_out[i]))
            num_len = str(index)
            out_file.write(";" + num_len)
            if 10 <= outer < 20:
                out_file.write(";" + str(11))
            elif 20 == outer:
                out_file.write(";" + str(10))
            else:
                out_file.write(";" + str(outer))
            out_file.write("\n")
            # 当分类为较短无效轨迹时，同时获取正序逆序序列
            if outer == 20:
                re_angle_out.reverse()
                distance_out.reverse()
                for i in range(len(re_angle_out)):
                    out_file.write(" " + str(re_angle_out[i]))
                out_file.write(";")
                for i in range(len(distance_out)):
                    out_file.write(" " + str(distance_out[i]))
                num_len = str(index)
                out_file.write(";" + num_len)
                out_file.write(";" + str(10))
                out_file.write("\n")
            out_file.close()
            inner += 1
            cap.release()
        outer += 1
    cv2.destroyAllWindows()
